#include <iostream>
#include <vector>
#include <string>
#include <stdlib.h>
#include <math.h>
#include <limits>
#include <fstream>
#include "fitness.hpp"
#include <sstream>
#include <algorithm>

using namespace std;

vector< vector<bool> > pool;
vector< float > eff; // selection coefficient
vector<string> refType;
vector<string> varType;

// generate poisson random numbers
inline int rpois(float lambda)
{
    float l=exp(-lambda);
    int k=0;
    float p=1.0;
    do
    {
        k++;
        p = p * ((float)rand()) / RAND_MAX;
    }
    while (p > l);
    return k-1;
}

inline double random(double start, double end) // cannot return exactly 0, or the program will fail
{
    double rndDbl = rand();
    rndDbl = double(start) + ((double)(end-start))*(((double)rndDbl)/((double)(RAND_MAX)));
    while(rndDbl == start ) {
    	rndDbl = rand();
		rndDbl = double(start) + ((double)(end-start))*(((double)rndDbl)/((double)(RAND_MAX)));
    }
    return double(rndDbl);
}

// cross two chromosomes, the actual number of recombinations happened is
// min + poisson(avg-min)
int sporulate(int min, float avg, vector<bool> &id1, vector<bool> &id2);

void randOneCross(vector<bool> &id1, vector<bool> &id2);

// the file named "fn" is no header, tab-delimited file with all the genome-types in it
void loadPool(string fn, int numMarker, int numChr, vector< vector<bool> > &res);

// load the selection coefficient "s".
// positive value means reference state is beneficial.
// negative value means variant state is beneficial.
void loadEffect(string fn, int numMarker, vector<float> &res);

// Get the frequencies of the reference state
void getFreq(vector< vector<bool> > &pop, int numMarker, int numChr, vector<float> &res);

void usage();

void outputPop(vector< vector<bool> > &pop, int numMarker, int numChr, string filename);

void outputFreq(vector<float> &freqList, int numMarker, string filename);

template <typename T>
string numberToString ( T num )
{
	stringstream ss;
	ss << num;
	return ss.str();
}

int main(int argc, char *argv[])
{
    if(argc != 19) usage();
    string file = argv[1]; // input file with all the markers in all the strains
    int nMarker = atoi(argv[2]); // number of markers
    int nHap = atoi(argv[3]); // number of chromosomes
    int minCross = atoi(argv[4]); // minimum number of cross over per generation
    float avgCross = atof(argv[5]); // average number of cross over per generation
    string effFile = argv[6]; // selection coefficient for the alternative alleles at each locus

    int nNormGen = atoi(argv[7]); // number of normal generations ,must be at least 3.
	// Note that in the experimental design, f1 have gone through no recombination, and f2 is the first generation after recombination
	// In this program, the generation number is based on the number of recombination the cells have gone through
    int normPopSize = atoi(argv[8]); // population size during normal generation, must be at least (nHap ^ 2)
	int seedNorm = atoi(argv[9]); // random seed for normal generation
	string outputNormAt = argv[10]; // when to output normal generation genotypes, "end", "all" or "none"
	string outputNorm = argv[11]; // filename to output normal generation genotypes
	int nPickNorm = atoi(argv[12]); // number of diploid to output genotype during normal generation

	int nSelectGen = atoi(argv[13]); // number of generation for selection generation
    int selectPopSize = atoi(argv[14]); // population size during selection generation
    int seedSel = atoi(argv[15]); // random seed for selection generation
    string outputFnFit = argv[16]; // filename to output the fitness
    string outputFnGt = argv[17]; // filename to output the genotypes
    string outputFnFreq = argv[18]; // filename to output the frequency of each strain after selection

    float sumFitness;
    float sumFitnessCopy;
    //float thisFitness;
    vector<float> allFitness;
    vector<float> tmpFitness;
    vector< vector<bool> > tmpPool;
    vector<float> selFreq;
    vector<int> allId;
    vector<int> tmpAllId;
    int curInd,tmpInt;
    vector<int> randPerm;

    int curNeed,curGot,curIdx;
    float x,curFit,rndDbl;

    loadPool(file,nMarker,nHap,pool);
    loadEffect(effFile,nMarker,eff);

	normPopSize = (normPopSize>pow(nHap,2))?normPopSize:pow(nHap,2);
	nNormGen = (nNormGen>2)?nNormGen:3;
    pool.reserve(normPopSize * 2);
    allFitness.reserve(normPopSize);
    tmpPool.reserve(normPopSize * 2);
    tmpFitness.reserve(normPopSize);
    tmpPool = pool;

    // Generations without selection, with mate and sporulation
    srand(seedNorm);
    // F1 -> F2 in real experiment. F1 as in the program's concept (first recombination)
    // Pairwise mating, no homozygote
    int indPerParent = int(normPopSize / nHap);
    for(int i=0;i<nHap;i++)
    {
        for(int j=0;j<indPerParent;j++)
        {
            int k = int(random(0.0,nHap));
            while( (i == k) or (k<=0) or (k>=nHap) ) { // no homozygote
                k = int(random(0.0,nHap));
            }
            if((signed)tmpPool.size() > (i * indPerParent + j) * 2 )
            {
                tmpPool[(i * indPerParent + j) * 2] = pool[i];
            }
            else
            {
                tmpPool.push_back(pool[i]);
            }
            if((signed)tmpPool.size() > (i * indPerParent + j) * 2 + 1 )
            {
                tmpPool[(i * indPerParent + j) * 2 + 1] = pool[k];
            }
            else
            {
                tmpPool.push_back(pool[k]);
            }

            sporulate(minCross,avgCross,tmpPool[(i * indPerParent + j) * 2],tmpPool[(i * indPerParent + j) * 2 + 1]); // first recombination
        }
    }
    pool = tmpPool;

    // F2 -> F3 in real experiment. F2 as in the program's concept (second recombination)
    // Mating happens only among those share the same parent
    for(int i=0;i<nHap;i++)
    {
        for(int j=0;j<indPerParent;j++)
        {
            int k = int(random(0.0,indPerParent * 2));
            int l = int(random(0.0,indPerParent * 2));
            if((signed)tmpPool.size() > (i * indPerParent + j) * 2 ) tmpPool[(i * indPerParent + j) * 2] = pool[i * indPerParent * 2 + k];
            else tmpPool.push_back(pool[i * indPerParent * 2 + k]);
            if((signed)tmpPool.size() > (i * indPerParent + j) * 2 + 1 ) tmpPool[(i * indPerParent + j) * 2 + 1] = pool[i * indPerParent * 2 + l];
            else tmpPool.push_back(pool[i * indPerParent * 2 + l]);
            sporulate(minCross,avgCross,tmpPool[(i * indPerParent + j) * 2],tmpPool[(i * indPerParent + j) * 2 + 1]); // second recombination
        }
    }
	pool = tmpPool;

	// remaining normal generations
    for(int i=2;i<nNormGen;i++)
    {
    	cerr << "Processing normal F" << i << endl;
    	if(outputNormAt.compare("all") == 0)
    	{
			string tmpFilename = "";
			tmpFilename.append(outputNorm);
			tmpFilename.append(".f");
			tmpFilename.append(numberToString(i));
			outputPop(pool, nMarker, nPickNorm * 2, tmpFilename);
			tmpFilename = "";
    	}
        for(int j=0;j<normPopSize;j++)
        {
            for(int k=0;k<2;k++)
            {
                tmpInt = rand() % pool.size();
                if((signed)tmpPool.size() > (j * 2 + k) )
                {
                    tmpPool[j * 2 + k] = pool[tmpInt];
                }
                else
                {
                    tmpPool.push_back(pool[tmpInt]);
                }
            }
            sporulate(minCross,avgCross,tmpPool[j*2],tmpPool[j*2 + 1]);
        }
        pool = tmpPool;
    }
	if((outputNormAt.compare("all") == 0) || (outputNormAt.compare("end") == 0))
	{
		string tmpFilename = "";
		tmpFilename.append(outputNorm);
		tmpFilename.append(".f");
		tmpFilename.append(numberToString(nNormGen));
		outputPop(pool, nMarker, nPickNorm * 2, tmpFilename);
		tmpFilename = "";
	}

    //cerr << "allId.clear()" << endl;
	allId.clear();
	//cerr << "selFreq.clear()" << endl;
	selFreq.clear();
	//cerr << "allFitness.clear()" << endl;
	allFitness.clear();
	for(int i=0;i<nPickNorm;i++)
	{
	    if((signed)tmpPool.size() > i) tmpPool[i] = pool[i * 2]; // don't pick from the sister gametes
	    else tmpPool.push_back(pool[i * 2]); // don't pick from the sister gametes
		allId.push_back(i); // id each haplotype
		selFreq.push_back(0); // get ready for final freqency calculation
		allFitness.push_back(getHapIntFitness(pool[i * 2],eff));
		// cerr << allFitness[allFitness.size() - 1];
	}
	pool = tmpPool;

	outputPop(pool, nMarker, nPickNorm, outputFnGt);
	outputFreq(allFitness, nPickNorm, outputFnFit);

    // Generations with selection
    srand(seedSel);
    for(int i=0;i<nSelectGen;i++) // each generation
    {
        cerr << "Selection Generation "<<i<<endl;
        sumFitness = 0.0;
        randPerm.clear();
        if(i == 0)
        {
            for(int j=0;j<nPickNorm;j++)
            {
                sumFitness += allFitness[j];
                randPerm.push_back(j);
            }
        }
        else
        {
            for(int j=0;j<selectPopSize;j++)
            {
                sumFitness += allFitness[j];
                randPerm.push_back(j);
            }
        }
        std::random_shuffle(randPerm.begin(),randPerm.end());
        //cerr << "allfitness.size()=" << allFitness.size() << endl;
        //cerr << "pool.size()=" << pool.size() << endl;

        //Selection
        // this selection prcedure is O(m + n)
        // start O(m + n)
		pool.reserve(selectPopSize);
		allFitness.reserve(selectPopSize);
		tmpPool.reserve(selectPopSize);
		tmpFitness.reserve(selectPopSize);
		allId.reserve(selectPopSize);
		tmpAllId.reserve(selectPopSize);
		curIdx = 0;
        curInd=randPerm[curIdx]; // Use a random permutation instead of the fixed order, so nobody will remain the end of the queue
        //if(i==0) curNeed = selectPopSize * 0.95; // Be conservative, so that the last few individual could have their chance
        //else curNeed = selectPopSize;
        // Should not need to do these, but maybe due to precision problem, the last few individuals are alway
        // under-sampled. And such bias is even worse by the fact that the random samples are intrinsically sorted.
        // Making the last sample(s) disappear as generation number increases.
        // Temporarily reducing the "need" reduces eveyone's sampling probability to the same degree.
        // And then we use the old fasion Rollette-Wheel method to fill in the blanks.
        curNeed = selectPopSize;
        curGot=0;
        curFit=allFitness[curInd];
        sumFitnessCopy = sumFitness;

        while(curNeed>0) {
            x = sumFitnessCopy * ( 1.0 - pow((random(0.0,1.0)),1.0/((float)curNeed) ) );
            sumFitnessCopy -= x;
            while(x > curFit) {
                x -= curFit;
                if( ( (i==0) && (curIdx >= nPickNorm - 1) ) || ( (i > 0) && (curIdx >= selectPopSize - 1) ) )
                {
                	break;
                }
                curIdx ++;
                curInd = randPerm[curIdx];
                curFit = allFitness[curInd];
            }
            curFit -= x;

            if((signed)tmpFitness.size() > curGot) tmpFitness[curGot] = allFitness[curInd];
            else tmpFitness.push_back( allFitness[curInd] );
            if((signed)tmpAllId.size() > curGot) tmpAllId[curGot] = allId[curInd];
            else tmpAllId.push_back( allId[curInd] );

            curGot ++;
            curNeed --;
            if(curGot % int(selectPopSize / 10) == 0) cerr << curGot << endl;
        }
        // end O(m + n)
        curNeed = selectPopSize - curGot;
        //cerr << curNeed << " short" << endl;
        if(curNeed > 0)
        {   // if the fast way did not get enough sample
            // fall back on the slow way
            for(int j=0;j<curNeed;j++)
            {
                rndDbl = random(0,sumFitness);
                curInd = -1;
                while((rndDbl >= 0) && (curInd < (signed)allFitness.size()-1) )
                {
                    curInd ++;
                    rndDbl = rndDbl - allFitness[curInd];
                }
                if((signed)tmpFitness.size() > curGot + j) tmpFitness[curGot + j] = allFitness[curInd];
                else tmpFitness.push_back( allFitness[curInd] );
                if((signed)tmpAllId.size() > curGot + j) tmpAllId[curGot + j] = allId[curInd];
                else tmpAllId.push_back( allId[curInd] );

                if((curGot + j + 1) % int(selectPopSize / 10) == 0) cerr << (curGot+j + 1) << endl;
            }
        }
        allFitness = tmpFitness;
        allId = tmpAllId;
    }

    // calculate the frequencies
    // getFreq(pool,nMarker ,selectPopSize * 2, selFreq);
    // used in xQTL, but not GWAS

	// calculate the freuqencies of each haplotype
	for(unsigned int i=0;i<allId.size();i++) selFreq[allId[i]] ++;
	for(int i=0;i<nPickNorm;i++)
	{
	    //cerr << "selFreq[i]" << selFreq[i] << endl;
	    selFreq[i] = selFreq[i] / allId.size();
	}
	//cerr << "allId.size()" << allId.size() << endl;

    // print results
	outputFreq(selFreq, nPickNorm, outputFnFreq);

    // output detail genome types
    // outputPop(tmpPool, nMarker, selectPopSize * 2, outputFn2);

    pool.clear();
    tmpPool.clear();
    allFitness.clear();
    tmpFitness.clear();
    allId.clear();
    tmpAllId.clear();
    exit(0);
}

void outputFreq(vector<float> &freqList, int numMarker, string filename)
{
	ofstream ofs;
    ofs.open(filename.c_str());
    for(int i = 0;i<numMarker;i++)
    {
        ofs << freqList[i] << endl;
    }
    ofs.close();
    return;
}

void outputPop(vector< vector<bool> > &pop, int numMarker, int numChr, string filename)
{
    ofstream ofs;
    ofs.open(filename.c_str());
    for(int i=0;i<numChr;i++ )
    {
        for(int j=0;j<numMarker;j++)
        {
            if(j == numMarker - 1)
            {
                if(pop[i][j]) {ofs << varType[j] << endl;}
                else{ ofs << refType[j] << endl;}
            }
            else
            {
                if(pop[i][j]) {ofs << varType[j] << '\t';}
                else{ ofs << refType[j] << '\t';}
            }
        }
    }
    ofs.close();
    return;
}

void getFreq(vector< vector<bool> > &pop, int numMarker, int numChr, vector<float> &res)
{
    int tmpInt;
    res.clear();
    for(int i=0;i<numMarker;i++)
    {
        tmpInt = 0;
        for(int j=0;j<numChr;j++)
        {
            if(!pop[j][i]) tmpInt ++;
        }
        //cerr << "Marker "<<i<<" : " << tmpInt << endl;
        res.push_back( ((float)tmpInt) / ((float)numChr) );
    }
    return;
}

void loadEffect(string fn, int numMarker, vector<float> &res)
{
    ifstream ifs;
    ifs.open(fn.c_str());
    float tmpDbl;
    res.clear();

    for(int i=0;i<numMarker;i++)
    {
        ifs >> tmpDbl;
        res.push_back(tmpDbl);
    }
    ifs.close();
    return;
}

void loadPool(string fn, int numMarker, int numChr, vector< vector<bool> > &res)
{
    int loadedCnt = 0;
    int markerCnt = 0;
    vector<bool> tmpVecBln;
    tmpVecBln.reserve(numMarker);
    char ch;
    ifstream ifs;
    ifs.open(fn.c_str());
    while(true)
    {
        ch = ifs.peek();
        string strCh(1,ch);
        switch(ch)
        {
        case '\t':
            if(numMarker == markerCnt)
            { // ignore the whole line
                ifs.ignore(numeric_limits<streamsize>::max(),'\n');
                if(loadedCnt % 100000 == 0) {cerr << "Loaded " << loadedCnt << " ; nMarker " << tmpVecBln.size() << endl;}
                res.push_back(tmpVecBln);
                loadedCnt ++;
                tmpVecBln.clear();
                markerCnt = 0;
                if(loadedCnt == numChr)
                {
                    ifs.close();
                    return;
                }
                continue;
            }
            break;
        case -1:
            if(markerCnt == numMarker)
            {
                if(loadedCnt % 100000 == 0) {cerr << "Loaded " << loadedCnt << " ; nMarker " << tmpVecBln.size() << endl;}
                res.push_back(tmpVecBln);
                loadedCnt ++;
                tmpVecBln.clear();
                markerCnt = 0;
                if(loadedCnt == numChr)
                {
                    ifs.close();
                    return;
                }
                else
                {
                    cerr << "Error loading data: the file contains "<< loadedCnt << " useful rows,"<< endl;
                    cerr << "whereas you requested " << numChr << "." << endl;
                    exit(1);
                }
            }
            else
            {
                cerr << "Error loading data: a row in the file contains "<<markerCnt<<" markers,"<<endl;
                cerr << "whereas you requested " << numMarker << "." << endl;
                exit(1);
            }
            break;
        case '\n':
            if(markerCnt == numMarker )
            {
                if(loadedCnt % 100000 == 0) {cerr << "Loaded " << loadedCnt << " ; nMarker " << tmpVecBln.size() << endl;}
                res.push_back(tmpVecBln);
                loadedCnt ++;
                tmpVecBln.clear();
                markerCnt = 0;
                if(loadedCnt == numChr)
                {
                    ifs.close();
                    return;
                }
            }
            else if(markerCnt < numMarker)
            {
                cerr << "Error loading data: a row in the file contains "<<markerCnt<<" markers,"<<endl;
                cerr << "whereas you requested " << numMarker << "." << endl;
                exit(1);
            }
            if(loadedCnt == numChr)
            {
                ifs.close();
                return;
            }
            break;
        default:
            if(loadedCnt == 0)
            {
                refType.push_back(strCh);
                varType.push_back(strCh);
                tmpVecBln.push_back(false);
            }
            else
            {
                if(refType[markerCnt].compare(strCh) == 0 )
                {
                    tmpVecBln.push_back(false);
                }
                else
                {
                    tmpVecBln.push_back(true);
                    varType[markerCnt] = strCh;
                }
            }
            markerCnt ++;
            break;
        }
        ifs.get();
    }
    ifs.close();
    return;
}

int sporulate(int min, float avg, vector<bool> &id1, vector<bool> &id2)
{
    int nCross = min + rpois(avg - min);
    for(int i=0; i<nCross; i++) randOneCross(id1,id2);
    return nCross;
}

void randOneCross(vector<bool> &id1, vector<bool> &id2)
{
    int at = 1 + (rand() % (id1.size() - 1));
    bool tmp;
    for(unsigned int i=at; i<id1.size(); i++)
    {
        tmp = id1[i];
        id1[i] = id2[i];
        id2[i] = tmp;
    }
    return;
}



void usage()
{
    cerr << "Expecting parameters are : (All are needed and in this paricular order)" << endl;
    /*
    cerr << "1.markerFile: A tab delimited file listing the markers' states in starting strains" << endl;
    cerr << "    Please use only one character to indicate the state of it." << endl;
    cerr << "    The states of the first row in this file will be regarded as reference state," << endl;
    cerr << "    and all other states are regarded as variants. (So >= 2 states will be reduced to 2)" << endl;
    cerr << "2.nMarker: How many markers to load." << endl;
    cerr << "    Must be smaller than number of columns in markerFile." << endl;
    cerr << "3.nChromosome: How many chromosomes to load." << endl;
    cerr << "    Must be smaller than number of rows in markerFile" << endl;
    cerr << "4.popSize: Population size." << endl;
    cerr << "5.minCross: Minimum number of crossover events per sporulation." << endl;
    cerr << "6.avgCross: Expected number of crossover events per sporulation." << endl;
    cerr << "    The actuall number of crossovers will be (minCross + poisson(avgCross-minCross))" << endl;
    cerr << "7.selCoefFile: A tab delimited file listing the selection coefficient for each marker." << endl;
    cerr << "    A positive value will favored the reference state, and negative value will favored the variant state" << endl;
    cerr << "    E.g. -0.05 means a relative fitness of 1.05 for variant state." << endl;
    cerr << "8.nGenNoSel: Number of generations without selection. During which population size is popSize" << endl;
    cerr << "    Mating and sporulation happend during these generations" << endl;
    cerr << "9.nGenSel: Number of generations with selection. During which population size is popSize." << endl;
    cerr << "    No mating and sporulation during these generations." << endl;
    cerr << "10.randSeed1: Random number seed for generations WITHOUT selection." << endl;
    cerr << "11.randSeed2: Random number seed for generations WITH selection." << endl;
    cerr << "12.freqOutputFile: Output file for the reference state frequencies at the end." << endl;
//    cerr << "13.popOutputFile: Output file for the whole population at the end." << endl;
	*/
    exit(0);
}
